package cn.doitedu.profile.ml.gender

import org.apache.spark.ml.classification.{LogisticRegression, NaiveBayes}
import org.apache.spark.ml.linalg
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.SparkSession

/**
 * @author 涛哥
 * @nick_name "deep as the sea"
 * @contact qq:657270652 wx:doit_edu
 * @site www.doitedu.cn
 * @date 2021-06-24
 * @desc 行为性别预测
 */
object ActionGenderPredict {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("")
      .master("local")
      .config("spark.sql.shuffle.partitioins", "1")
      .getOrCreate()

    import spark.implicits._

    // 加载训样本数据
    val sample1 = spark.read.options(Map("header" -> "true", "inferSchema" -> "true")).csv("profile/data/gender/sample")
    val sample2 = spark.read.options(Map("header" -> "true", "inferSchema" -> "true")).csv("profile/data/gender/test")


    // 向量化
    val sample = sample1.union(sample2)

    val vecDf = sample.map(row=>{
      val label = row.getAs[Double]("label")
      val guid = row.getAs[Double]("guid")

      val features1 = Array(
        row.getAs[Double]("category1") ,
        row.getAs[Double]("category2") ,
        row.getAs[Double]("category3") ,
        row.getAs[Double]("brand1")    ,
        row.getAs[Double]("brand2")    ,
        row.getAs[Double]("brand3")
      )


      val features2 = Array(
        row.getAs[Double]("day30_buy_cnts"),
        row.getAs[Double]("day30_buy_amt")
      )


      (label,guid,Vectors.dense(features1),Vectors.dense(features2))
    }).toDF("label","guid","features1","features2")


    // 将样本划分成训练集和测试集
    val arr = vecDf.randomSplit(Array(0.8, 0.2))
    val train = arr(0)
    val test = arr(1)

    // 构造算法
    val bayes = new NaiveBayes()
      .setLabelCol("label")
      .setFeaturesCol("features1")
      .setSmoothing(1.0)


    val regression = new LogisticRegression()
      .setLabelCol("label")
      .setFeaturesCol("features2")
      .setRegParam(0.1)

    // 训练模型
    val bayesModel = bayes.fit(train)
    val regressionModel = regression.fit(train)


    // 模型预测
    val bayesPrediction = bayesModel.transform(test)
    val regressionPrediction = regressionModel.transform(test)

    val vecExtract = (vec:linalg.Vector)=>{
      vec.toArray
    }
    spark.udf.register("vec_extract",vecExtract)


    bayesPrediction.selectExpr("label","guid","vec_extract(probability)[0] as prob0","vec_extract(probability)[1] as prob1").createTempView("bayes")
    regressionPrediction.selectExpr("label","guid","vec_extract(probability)[0] as prob0","vec_extract(probability)[1] as prob1").createTempView("regr")

    spark.sql(
      """
        |
        |select
        |bayes.label,
        |bayes.guid,
        |if(0.2* bayes.prob0 + 0.8*regr.prob0>0.5,0,1) as predict_gender
        |
        |from bayes join regr on bayes.guid=regr.guid
        |
        |""".stripMargin).show(100,false)


    spark.close()

  }

}
